from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
import ufl
from ufl import dx, ds, dS, grad, inner, dot, avg, jump
from dolfinx import mesh, fem, io, la
import basix
from dolfinx.fem.petsc import assemble_matrix, assemble_vector


from mpi4py import MPI
import numpy as np
from petsc4py import PETSc
import ufl
from ufl import dx, ds, dS, grad, inner, dot, avg, jump
from dolfinx import mesh, fem
from dolfinx.fem.petsc import assemble_matrix, assemble_vector, apply_lifting, set_bc
from dolfinx import io

# ------------------ 函数定义 ------------------
def u_exact_func(x, t):
    return np.exp(x[0] + x[1] + t)

def u_exact_expr(x, t):
    return ufl.exp(x[0] + x[1] + t)

def f_func(x, t):
    return -np.exp(x[0] + x[1] + t)

# ------------------ 参数 ------------------
type_ = "SIPG"            
epsilon = -1.0       
k = 2                     
sigma_0_base = 6
CFL = 0.02/(k**2)
x_min, x_max, y_min, y_max = 0.0, 2.0, 0.0, 1.0
T = 1.0

h_list = [1/4, 1/8, 1/16, 1/32, 1/64]
h1_semi_err, l2_err = [], []

for h in h_list:
    dt = h**2 / 2
    N = int(round(1 / h))
    domain = mesh.create_rectangle(
        MPI.COMM_WORLD,
        [[x_min, y_min], [x_max, y_max]],
        [int((x_max - x_min) / h), int((y_max - y_min) / h)],
        mesh.CellType.triangle,
    )

    x = ufl.SpatialCoordinate(domain)
    V = fem.functionspace(domain, ("DG", k))

    u = ufl.TrialFunction(V)
    v = ufl.TestFunction(V)

    n = ufl.FacetNormal(domain)
    h_cell = ufl.CellDiameter(domain)
    h_avg = ufl.avg(h_cell)

    sigma_int = sigma_0_base * (k + 1) * k / h_avg
    sigma_bnd = sigma_0_base * (k + 1) * k / h_cell

    # ------- DG离散化矩阵 A -------
    a_vol = inner(grad(u), grad(v)) * dx
    a_int = (
        - inner(avg(grad(v)), jump(u, n)) * dS
        + epsilon * inner(avg(grad(u)), jump(v, n)) * dS
        + sigma_int * inner(jump(u, n), jump(v, n)) * dS
    )
    a_D = (
        epsilon * dot(grad(u), n) * v * ds
        - dot(grad(v), n) * u * ds
        + sigma_bnd * u * v * ds
    )
    A_form = a_vol + a_int + a_D

    # 质量矩阵
    M_form = inner(u, v) * dx

    A = assemble_matrix(fem.form(A_form))
    A.assemble()
    M = assemble_matrix(fem.form(M_form))
    M.assemble()

    # 初值
    u_n = fem.Function(V)
    u_n.interpolate(lambda x: np.exp(x[0] + x[1] + 0.0))

    u_new = fem.Function(V)

    # 右端项形式
    f_expr = fem.Function(V)

    # 组装常规模型矩阵
    system_mat = M.copy()
    # axpy: A * X plus Y
    # dt是标量系数 A是加上的系数 false表示矩阵结构不变
    # system_mat = M + dt * A
    system_mat.axpy(dt, A, False)
    system_mat.assemble()

    ksp = PETSc.KSP().create(MPI.COMM_WORLD)
    ksp.setOperators(system_mat)
    ksp.setType("preonly")
    ksp.getPC().setType("lu")

    t = 0.0
    while t < T - 1e-10:
        t += dt
        # 更新 f
        f_expr.interpolate(lambda x: f_func(x, t))

        # 右端向量
        L_form = inner(f_expr, v) * dx \
            + epsilon * dot(grad(v), n) * u_exact_expr(x, t) * ds \
            + sigma_bnd * u_exact_expr(x, t) * v * ds
        L_vec = assemble_vector(fem.form(L_form))
        L_vec.assemble()

        # 添加 M*u_n
        b = M.createVecRight()    # 创建一个PETSc右手边向量（与M兼容）
        M.mult(u_n.x.petsc_vec, b)  # b = M * u_n
        b.axpy(dt, L_vec)          # b = b + dt * L_vec

        ksp.solve(b, u_new.x.petsc_vec)
        u_new.x.scatter_forward()

        # 存储旧时间步解
        u_n.x.array[:] = u_new.x.array[:]

    # -------- 误差计算 --------
    u_e = fem.Function(V)
    u_e.interpolate(lambda x: u_exact_func(x, T))

    diff = fem.Function(V)
    diff.x.array[:] = u_n.x.array[:] - u_e.x.array[:]

    L2_err = np.sqrt(fem.assemble_scalar(fem.form(inner(diff, diff) * dx)))
    H1_semi = np.sqrt(fem.assemble_scalar(fem.form(inner(grad(diff), grad(diff)) * dx)))

    l2_err.append(L2_err)
    h1_semi_err.append(H1_semi)

    if MPI.COMM_WORLD.rank == 0:
        print(f"h={h:.4f}, L2误差={L2_err:.3e}, H1半范数误差={H1_semi:.3e}")

# -------- 收敛阶统计 --------
def _rates(errs, hs):
    out = []
    for i in range(len(errs) - 1):
        out.append(np.log(errs[i] / errs[i+1]) / np.log(hs[i] / hs[i+1]))
    return out

l2_list = _rates(l2_err, h_list)
h1_semi_list = _rates(h1_semi_err, h_list)

if MPI.COMM_WORLD.rank == 0:
    print("\n===== 收敛结果 =====")
    print("h:               ", ["{:.5f}".format(h) for h in h_list])
    print("L2收敛阶:         ", ["{:.2f}".format(e) for e in l2_list])
    print("H1半范数收敛阶:   ", ["{:.2f}".format(e) for e in h1_semi_list])